import os
import uuid
from typing import List, Optional
import subprocess

from osgeo import gdal

from common.logger import logger
from common.typing_utils import numeric

from common.geometry.geom import Geom
from common.geometry.bbox import BoundingBox
from common.geoimage.gdal_sys import GDALRaster, warp_ds
from common.geoimage.raster_dataset import RasterDataset
from common.geoimage.virtual_raster import VRTScene, VRTSceneCollection

from catalog.models.base_stac_item import BaseSTACItem
from catalog.raster.utils import (
    determine_dest_srs,
    determine_dest_res,
    determine_dest_extent,
)

from catalog.raster.strategy_s2decloud_v1 import (
    get_s2_l2a_add_boa_offset,
    update_s2_l2a_boa,
)


gdal.UseExceptions()


def mosaic_geoimages(
    geoimages,
    output_fpath=None,
    x_res=None,
    y_res=None,
    dst_srs=None,
    output_bbox=None,
    fpath_aoi=None,
    crop2aoi=False,
    resample_alg="nearest",
    src_nodata=None,
    dst_nodata=None,
):
    """Mosaic multiple geo-images with the same spatial reference.

    Parameters
    ----------
    geoimages: list
        List of path strings or gdal.Datasets to multiple geo-images.
    output_fpath: str, optional
        Output path string to the mosaicked geo-image.
    x_res: float
        Pixel resolution in `x` axis.
    y_res: float
        Pixel resolution in `y` axis.
    output_srs: osr.SpatialReference, optional
        Output spatial reference system, if not specified, use the one from the first file.
        Must in the form of EPSG string, e.g., "EPSG:4326"
    output_bbox: BoundingBox
        Bounding box of the AOI for cropping, (left, bottom, right, upper)
    fpath_aoi: string
        Path to the AOI for cropping the mosaicked image.
    crop2aoi: bool
        Whether crop image to AOI.
    resample_alg: string
        Resampling method for image reprojection, e.g., 'bilinear', 'near', 'cubic', etc.
    src_nodata: int, optional
        Pixel value of no data in source geoimage.
    dst_nodata: int, optional
        Pixel value of no data in destination geoimage.

    Returns
    -------
    mosaicked_ds: gdal.Dataset object
        Mosaicked GDAL dataset object.
    """
    if not isinstance(geoimages, list):
        raise TypeError("Invalid type of 'geoimage_fpaths'. Accepts a list")

    if not geoimages:
        raise RuntimeError("Empty list of geoimage file paths.")

    resample_alg = "near" if resample_alg == "nearest" else resample_alg

    geoimage_dss = []
    for i, geoimage in enumerate(geoimages):
        vsimem_fpath = f"/vsimem/{uuid.uuid4()}.tif"
        ds = gdal.Warp(
            vsimem_fpath,
            geoimage,
            dstSRS=dst_srs,
            resampleAlg=resample_alg,
            srcNodata=src_nodata,
            dstNodata=dst_nodata,
            multithread=True,
            options=[
                "NUM_THREADS=8",
            ],
        )
        if gdal.VSIStatL(vsimem_fpath) is not None:
            gdal.Unlink(vsimem_fpath)
        if ds is None:
            raise Exception(f"GDAL failed to open or reproject: {geoimage}.")
        geoimage_dss.append(ds)

    if output_fpath is None:
        output_fpath = "/vsimem/tmp_output_geoimage.tif"
    try:
        if output_bbox is not None:
            output_bbox = (
                output_bbox.left,
                output_bbox.bottom,
                output_bbox.right,
                output_bbox.upper,
            )
        # if AOI is specified, ignore output bbox
        if output_bbox is not None and fpath_aoi is not None:
            output_bbox = None
        mosaicked_fpath = "/vsimem/tmp_mosaicked_geoimage.tif"
        mosaicked_ds = gdal.Warp(
            mosaicked_fpath,
            geoimage_dss,
            outputBounds=output_bbox,
            outputBoundsSRS=dst_srs if output_bbox is not None else None,
            xRes=x_res,
            yRes=y_res,
            resampleAlg=resample_alg,
            srcNodata=src_nodata,
            dstNodata=dst_nodata,
            cutlineDSName=fpath_aoi,
            cropToCutline=crop2aoi,
            multithread=True,
            options=[
                "NUM_THREADS=8",
            ],
        )
    except Exception as e:
        raise Exception(f"Failed to mosaic multiple geoimages due to err `{e}`.")
    else:
        if mosaicked_ds is None:
            raise Exception("GDAL failed to mosaic multiple geoimages.")
    finally:
        for i in range(len(geoimage_dss)):
            geoimage_dss[i] = None
        geoimage_dss = None
        if gdal.VSIStatL(mosaicked_fpath) is not None:
            gdal.Unlink(mosaicked_fpath)
        if gdal.VSIStatL(output_fpath) is not None:
            gdal.Unlink(output_fpath)

    return mosaicked_ds


def composite_recent_s2_v1(
    items: List[BaseSTACItem],
    asset_names: List[str],
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    resample_alg: str = "near",
):
    """Composite a given list of Sentinel-2 items using the show-most-recent-data strategy

    NOTE: This function is designed to be used within `catalog.raster.raster` only.
    It relies on the caller to provide several arguments correctly.
    """
    items = sorted(items, key=lambda item: item.start_datetime, reverse=True)
    logger.debug(str(items))

    try:
        nodata_values = []
        for asset_name in asset_names:
            nodata_values.append(items[0].bands[asset_name].nodata_val)
        if len(nodata_values) == 1:
            nodata_values = nodata_values[0]
    except Exception:
        logger.warn(
            "cannot read `nodata_values` from given product's default settings, fall back to use `0`."
        )
        nodata_values = 0

    # read all image data into memery to handle s2 offset issue
    loaded_rds = []
    for item in items:
        vrt_scene = item.choose_assets(asset_names, dest_nodata=nodata_values)
        if vrt_scene is None:
            return None
        rds = RasterDataset.from_vrt_scene(vrt_scene)
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            rds.data = update_s2_l2a_boa(rds.data, asset_offsets)
        loaded_rds.append(rds)

    if dest_srs is None:
        dest_srs = determine_dest_srs(aoi, items)

    if dest_res is None:
        dest_res = determine_dest_res(dest_srs, items)

    dest_extent = determine_dest_extent(extent, aoi, dest_srs, dest_res)

    # Convert all raster-dataset to gdal.Dataset
    dss = []
    for rds in loaded_rds:
        with rds.to_gdal_raster() as gr:
            dss.append(gr.ds)
        del rds
    del loaded_rds

    # Mosaic, warp, and crop
    mosaicked_ds = mosaic_geoimages(
        dss,
        x_res=dest_res,
        y_res=dest_res,
        dst_srs=dest_srs,
        output_bbox=dest_extent,
        fpath_aoi=None,
        crop2aoi=False,
        resample_alg=resample_alg,
        src_nodata=nodata_values,
        dst_nodata=nodata_values,
    )
    for ds in dss:
        del ds
    del dss
    with GDALRaster.wrap_gdal_dataset(mosaicked_ds) as gr:
        rds = RasterDataset.from_gdal_raster(gr)

    return rds


def composite_recent_s2_v2(
    items: List[BaseSTACItem],
    asset_names: List[str],
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    resample_alg: str = "near",
):
    """Composite a given list of Sentinel-2 items using the show-most-recent-data strategy

    NOTE: This function is designed to be used within `catalog.raster.raster` only.
    It relies on the caller to provide several arguments correctly.
    """
    items = sorted(items, key=lambda item: item.start_datetime, reverse=True)
    logger.debug(str(items))

    try:
        nodata_values = []
        for asset_name in asset_names:
            nodata_values.append(items[0].bands[asset_name].nodata_val)
        if len(nodata_values) == 1:
            nodata_values = nodata_values[0]
    except Exception:
        logger.warn(
            "cannot read `nodata_values` from given product's default settings, fall back to use `0`."
        )
        nodata_values = 0

    # edit the VRT of each scene to apply the offset
    vrt_scenes = []
    for item in items:
        vrt_scene = item.choose_assets(asset_names, dest_nodata=nodata_values)
        if vrt_scene is None:
            return None
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            # set offset in the vrt xml.
            cmds = [
                "gdal_edit.py",
                vrt_scene._vrt_file_path,
                "-offset",
            ] + [str(o) for o in asset_offsets.tolist()]
            ret = subprocess.run(cmds)
            if ret.returncode != 0:
                raise Exception(
                    f"Failed to add offset to VRT {vrt_scene._vrt_file_path}."
                )
            # apply offset to data
            dir_vrt = os.path.dirname(vrt_scene._vrt_file_path)
            vrt_fpath_unoffset = os.path.join(dir_vrt, f"{uuid.uuid4()}.vrt")
            try:
                ds = gdal.Translate(
                    vrt_fpath_unoffset,
                    vrt_scene._vrt_file_path,
                    options="-ot UInt16 -unscale",
                )
                if ds is None:
                    raise Exception(
                        f"Failed to apply offset to VRT {vrt_fpath_unoffset}"
                    )
                ds = None
            except Exception as e:
                ds = None
                raise Exception(
                    f"Failed to apply offset to VRT {vrt_fpath_unoffset} "
                    f"due to err `{e}`"
                )
            vrt_scene_unoffset = VRTScene.from_vrt_file_path(vrt_fpath_unoffset)
            vrt_scenes.append(vrt_scene_unoffset)
        else:
            vrt_scenes.append(vrt_scene)

    vrt_scene_collection = VRTSceneCollection(vrt_scenes)

    if dest_srs is None:
        dest_srs = determine_dest_srs(aoi, items)

    if dest_res is None:
        dest_res = determine_dest_res(dest_srs, items)

    dest_extent = determine_dest_extent(extent, aoi, dest_srs, dest_res)

    vrt_scene_collection_warped = vrt_scene_collection.warp_to_common_srs(
        dest_srs=dest_srs, resample_alg=resample_alg, dest_nodata=nodata_values
    )
    vrt_scene = vrt_scene_collection_warped.mosaic_vrt()
    with GDALRaster.open(vrt_scene._vrt_file_path) as gdal_rst:
        warped_ds = warp_ds(
            gdal_rst,
            dest_srs=dest_srs,
            dest_extent=dest_extent,
            dest_nodata=nodata_values,
            dest_res=dest_res,
            resample_alg=resample_alg,
        )
        with warped_ds:
            rds = RasterDataset.from_gdal_raster(warped_ds, extent=dest_extent)

    return rds
